/* Source.cpp version

BTP_RainRemoval distribution version 1


The program is free for academic use. Please contact Amit Kumar, Roll No-09EC3521, Department of E&ECE, IIT Kharagpur
<iitkgp.ece.amit@gmail.com> if you are interested in using the software for
commercial purposes.

The software must not be modified and distributed without prior
permission of the author.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/

#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <iostream>
#include <fstream>
#include <stdio.h>
#include <cmath>


using namespace cv;
using namespace std;

int main(int, char** argv)
{
    VideoCapture cap(argv[1]); // open the file
    if(!cap.isOpened())  // check if we succeeded
        return -1;

    /////////////////////////////////////// Variable Declaration /////////////////////////////////////////////////////

    Mat frame;
    Mat edges;
    Mat prev_frame;
    Mat cur_frame;
    Mat next_frame;
    Size image_size((int)cap.get(CV_CAP_PROP_FRAME_WIDTH),(int)cap.get(CV_CAP_PROP_FRAME_HEIGHT));
    Mat cur_temp=Mat::zeros(image_size,CV_8UC1);
    Mat cur_temp_smooth=Mat::zeros(image_size,CV_8UC1);
    Mat prev_temp=Mat::zeros(image_size,CV_8UC1);    
    Mat cur_edge=Mat::zeros(image_size,CV_8UC1);
    Mat prev_edge=Mat::zeros(image_size,CV_8UC1);
    Mat next_edge=Mat::zeros(image_size,CV_8UC1);
    Mat next_temp;
    Mat img_diff;
    Mat diff3 = Mat::zeros(image_size,CV_8UC1);
    Mat diff4 = Mat::zeros(image_size,CV_8UC1);  
    Mat mask = Mat::zeros(image_size,CV_8UC1); 
    Mat diff_staticedge = Mat::zeros(image_size,CV_8UC1); 
    Mat diff_moveedge = Mat::zeros(image_size,CV_8UC1); 
    Mat recon_image = Mat::zeros(image_size,CV_8UC1); 
    Mat range = Mat::zeros(image_size,CV_8UC1);
    Mat spreadsymm = Mat::zeros(image_size,CV_8UC1);
    Mat classify = Mat::zeros(image_size,CV_8UC1);
    vector<float> tpf,tnf;
    int max,min;
    float mean,T1,T2,below_std,above_std;
	clock_t t1,t2;
    float diff=0,time_frame=0;
	int no_frames=0;
    T1=40;
    T2=4;

    double fps=cap.get(CV_CAP_PROP_FPS);
    VideoWriter writer("output_video.avi",CV_FOURCC('M', 'J', 'P', 'G'),fps, image_size);


    /////////////////////////////////////////// File Opening and setting /////////////////////////////////////////////////////

	ofstream res_file;

    res_file.open("file.txt");

	if(res_file.is_open())
	{
        cout<<"File open.\n";
        res_file<<"Results:\n";
	}
	else	
		cout<<"Unable to open the file.\n";

    ///////////////////////******************************** Code **********************************/////////////////////////

    for( ; ; )
    {
		no_frames++;
        cap>>frame; 
        //imshow("prev_video", prev_temp);        
        imshow("current video", cur_temp);        
        cvtColor(frame, next_temp, CV_BGR2GRAY);
        //imshow("next_video", next_temp);

		t1=clock();
                                  
  //////////////////////////////////////////////////  candidate rain pixels ////////////////////////////////////////////////
        
        uchar* data_cur=(uchar*)cur_temp.data;
        uchar* data_prev=(uchar*)prev_temp.data;
        uchar* data_next=(uchar*)next_temp.data;
        uchar* data_diff3=(uchar*)diff3.data;
        int width_step=cur_temp.step1();
        for(int i=0;i<image_size.height;i++)
        {
          for(int j=0;j<image_size.width;j++)
           {
             int k_cur = data_cur[i*width_step+j];
             int k_pre = data_prev[i*width_step+j];
             int k_next = data_next[i*width_step+j];             
             if(((k_cur-k_pre)>=7) || ((k_cur-k_next)>=7))//|| (((2*k_cur-k_pre-k_next)==0)))
              {
                data_diff3[i*width_step+j]=255;
              } 
             else
                data_diff3[i*width_step+j]=0;                           
           }
         }
        
        imshow("diff", diff3);
        
    //////////////////////////////////////////// static edge detection and removal /////////////////////////////////////////
    
        boxFilter(cur_temp,cur_temp_smooth,8, Size(5,5));
        Canny(prev_temp,prev_edge,10,125,3);
        Canny(next_temp,next_edge,10,125,3);
        Canny(cur_temp,cur_edge,10,125,3);
        uchar* data_prev_edge=(uchar*)prev_edge.data;
        uchar* data_next_edge=(uchar*)next_edge.data;
        uchar* data_cur_edge=(uchar*)cur_edge.data;
        
        Canny(cur_temp_smooth, edges,10,125,3);
        
        uchar* data_diff_staticedge=(uchar*)diff_staticedge.data;
        uchar* data_edge=(uchar*)edges.data;
        
        for(int i=0;i<image_size.height;i++)
        {
          for(int j=0;j<image_size.width;j++)
           { 
           if((data_prev_edge[i*width_step+j]==255)&&(data_cur_edge[i*width_step+j]==255)&&(data_next_edge[i*width_step+j]==255))            
              data_diff_staticedge[i*width_step+j]=255;  
           else
              data_diff_staticedge[i*width_step+j]=0;               
           }
        }        
        
        for(int i=0;i<image_size.height;i++)
        {
          for(int j=0;j<image_size.width;j++)
           { 
             
                data_diff3[i*width_step+j]=data_diff3[i*width_step+j]-data_diff_staticedge[i*width_step+j];
                                       
           }
        }
        
        std::vector<std::vector<Point> > contour;
        findContours(diff3,contour,CV_RETR_LIST,CV_CHAIN_APPROX_NONE);
        Mat cont = Mat::zeros(image_size,CV_8UC3);
        std::vector<std::vector<Point> > contour_boundary;
        
        for(int k=0;k<contour.size();k++)
        {
          double area = contourArea(contour[k]);
          if(area>=8 && area<=40)
          {
            contour_boundary.push_back(contour[k]);            
          }
        } 
        
        drawContours(cont,contour_boundary,-1,Scalar(255,255,255),-1);
        
        std::vector<Mat> cont_filled;
        split(cont,cont_filled);        
        
        diff3=diff3-cont_filled[0];
        //imshow("diff_refined",diff3);        
        
        
   ////////////////////////////////////////////////// moving edge detection ////////////////////////////////////////////     
        
        cur_edge.copyTo(diff_moveedge);
        uchar* data_diff_moveedge=(uchar*)diff_moveedge.data;
        
        for(int i=0;i<image_size.height;i++)
        {
            for(int j=0;j<image_size.width;j++)
            {
                 int temp=data_cur_edge[i*width_step+j];
                 int temp_pre=data_prev_edge[i*width_step+j];
                 int temp_next=data_next_edge[i*width_step+j];
                 if ((temp==temp_pre) && (temp==temp_next))
                 {
                    data_diff_moveedge[i*width_step+j]=0;
                 }
            }
        }
        
        diff3=diff3-diff_moveedge; 
        
        diff4=next_temp-diff3;
                 
        imshow("diff_moveedge",diff3);
        
        imshow("Only edges",diff4);
        
   //////////////////////////////////////////////////// Feature Extraction ////////////////////////////////////////////////
  

   uchar* data_range=(uchar*)range.data;
   
   for(int i=0;i<image_size.height;i++)
   {
    for(int j=0;j<image_size.width;j++)
     {
       max= data_cur[i*width_step+j];
       min= data_cur[i*width_step+j];
       if(max < data_next[i*width_step+j] )
       {
        max = data_next[i*width_step+j];
       }
       else if ( max < data_prev[i*width_step+j])
       {
        max = data_prev[i*width_step+j];
       }
       if (min > data_next[i*width_step+j])
       {
        min = data_next[i*width_step+j];
       }
       else if (min > data_prev[i*width_step+j])
       {
        data_next[i*width_step+j];
       }
       
       data_range[i*width_step+j]= max-min;        
     }
   }

   uchar* data_spreadsymm = (uchar*)spreadsymm.data;
   
   for(int i=0;i<image_size.height;i++)
   {
    for(int j=0;j<image_size.width;j++)
     {
       mean = (data_cur[i*width_step+j]+data_next[i*width_step+j]+data_cur[i*width_step+j])/3.0;

       if(data_prev[i*width_step+j]<mean)
       {
           if(data_cur[i*width_step+j]<mean)
           {
               above_std=0;
               below_std=(float)sqrt(pow((data_prev[i*width_step+j]-mean),2)+pow((data_cur[i*width_step+j]-mean),2))/2.0;
           }
       }

       else if (data_cur[i*width_step+j]>mean)
       {
           if(data_next[i*width_step+j]>mean)
           {
               below_std=0;
               above_std=(float)sqrt(pow((data_cur[i*width_step+j]-mean),2)+pow((data_next[i*width_step+j]-mean),2))/2.0;
           }
       }

       else if(data_cur[i*width_step+j]<mean)
       {
           if(data_next[i*width_step+j]>mean)
           {
               below_std=0;
               above_std=(float)sqrt(pow((data_prev[i*width_step+j]-mean),2)+pow((data_next[i*width_step+j]-mean),2))/2.0;
           }
       }

       else if(data_next[i*width_step+j]<mean)
       {
           if(data_prev[i*width_step+j]>mean)
           {
               below_std=0;
               above_std=(float)sqrt(pow((data_prev[i*width_step+j]-mean),2)+pow((data_cur[i*width_step+j]-mean),2))/2.0;
           }
       }
       else if(data_cur[i*width_step+j]<mean)
       {
           if(data_next[i*width_step+j]<mean)
           {
               above_std=0;
               below_std=(float)sqrt(pow((data_cur[i*width_step+j]-mean),2)+pow((data_next[i*width_step+j]-mean),2))/2.0;
           }
       }
       else if(data_next[i*width_step+j]<mean)
       {
           if(data_prev[i*width_step+j]<mean)
           {
               above_std=0;
               below_std=(float)sqrt(pow((data_prev[i*width_step+j]-mean),2)+pow((data_next[i*width_step+j]-mean),2))/2.0;
           }
       }

       data_spreadsymm[i*width_step+j]=abs(above_std-below_std);

     }
   }  

   
   //////////////////////////////////////////////////// Classification ////////////////////////////////////////////////////

   uchar* data_classify = (uchar*)classify.data;
   
   for(int i=0;i<image_size.height;i++)
   {
    for(int j=0;j<image_size.width;j++)
    {
      if( (data_range[i*width_step+j]>60) && (data_spreadsymm[i*width_step+j]>4) )
        data_classify[i*width_step+j]=255;
    }    
   }

   imshow("Classifier",classify);


   //////////////////////////////////////////////////// Parameter Calculation /////////////////////////////////////////////
   
   /*********************************************** True Positives and False Negatives ************************************/
   /*
   int count_tp,int count tn;
   
   for(int i=0;i<image_size.height;i++)
   {
    for(int j=0;j<image_size.width;j++)
    {
      if( (data_classify[i*width_step+j]==) && ())
      { 
        count_tp++; 
        tpf.push_back(count_tp++);
      }
      else 
      {
      count_tn++;
        tnf_pushback(count_tn);
      }
    }    
   }    
   
   */
   //////////////////////////////////////////////////// Image Inpainting /////////////////////////////////////////////////
        
        uchar* data_reconimg=(uchar*)recon_image.data;
        
        for(int p=0;p<image_size.height;p++)
        {
            for(int q=0;q<image_size.width;q++)
            {
              if(data_diff3[p*width_step+q])
              { 
                data_reconimg[p*width_step+q]=(data_prev[p*width_step+q]+data_next[p*width_step+q])/2;                
              }
              else
              {
                data_reconimg[p*width_step+q]=data_cur[p*width_step+q];
              }               
            }
        }

        GaussianBlur(recon_image,recon_image,Size(1,1),1.5,1.5);

  ////////////////////////////////////////////// Writing time in file ////////////////////////////////////////////////////
		t2=clock();

        diff=float(t2-t1)/CLOCKS_PER_SEC;

        time_frame+= diff;

        res_file<<"Time for frame "<<no_frames<<" is: "<<diff<<endl;

        imshow("Reconstructed Image",recon_image);

   ////////////////////////////////////////////////// writing video file /////////////////////////////////////////////////

        Mat color = Mat::zeros(image_size,CV_8UC3);
        Mat gray[] = {recon_image,recon_image,recon_image};
        merge(gray,3,color);
        writer.write(color);
        
   ///////////////////////////////////////////////// Taking next frame ///////////////////////////////////////////////////     
        

        cur_temp.copyTo(prev_temp);
        next_temp.copyTo(cur_temp);
                      
        if(waitKey(200) >= 0) break;
    }

    //cvReleaseVideoWriter(&writer);
    cap.release();

    res_file<<"Average Time per frame:"<<time_frame/no_frames<<endl;
    res_file.close();
	
    return 0;
}
